#!/usr/bin/python
##    Copyright (c) 2006 Greg Robinson. Email: gregarobi@dev.java.net
##
##    Permission is hereby granted, free of charge, to any person obtaining a copy of this
##    software and associated documentation files (the "Software"), to deal in the Software
##    without restriction, including without limitation the rights to use, copy, modify,merge,
##    publish, distribute, sublicense, and/or sell copies of the Software, and to permit
##    persons to whom the Software is furnished to do so, subject to the following conditions:
##
##    The above copyright notice and this permission notice shall be included in all copies
##    or substantial portions of the Software.
##
##    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
##    INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
##    PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
##    HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF
##    CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR
##    THE USE OR OTHER DEALINGS IN THE SOFTWARE.

import math, numpy, random, time


#===================================
##    Swarm creates swarm of "nParticles" particles in "nDimensions"-dimensional hyperspace
##    with dimension lengths "spaceSize". Each particle has a position vector contained in
##    "self.position" and a corresponding velocity vector contained in "self.velocity".
##    
##    You can think of this as a swarm of bees, where each bee is adjusting its velocity
##    towards the "best" bee in the swarm. The "best" bee is the one that, for example,
##    found the most honey.
##
##    This optimization can be applied to any smooth and continuous function of any number
##    of variables. "Swarm.fitnessFunc()" should be replaced with the applicable function
##    by way of inheritence.
##
##    TODO:
##        -Convergence instead of absolute error.
##            Sometimes, and often, the optimum value is not known. Add functionality for
##            to find relative convergence of the swarm.
##        -Eliminate hard-coded maximum fitness currently on updateFitness. There are times
##            when less is better or when one is trying to find a specific value.
##        -Neighborhood swarming. "self.metric" is currently unused, but will be used to find
##            neighbor particles. This functionality is currently in progress. It should solve
##            most of the local optima difficulties this currently has.
##        -More flexible algorithm for updateVelocity().
##            Maybe add more weight options such as "crazyness"?
##        -Faster method for updateVelocity natively using numpy?
##            According to unittest, this is one of the slowest parts.
##        -Comments! I got carried away with coding, so it might be difficult for
##            other people to understand it.
class Swarm:
    def __init__(self, nDimensions, spaceSize, nParticles, learn_rate):
        self.LEARN_RATE = learn_rate
        self.nDimensions = nDimensions
        self.spaceSize = spaceSize
        self.nParticles = nParticles
        self.position = numpy.random.random((nParticles,nDimensions)) * spaceSize
        self.MAX_POSITION = numpy.ones(nDimensions) * spaceSize
        self.MIN_POSITION = numpy.zeros(nDimensions)
        self.velocity = numpy.random.uniform(-spaceSize, spaceSize, (nParticles,nDimensions))
        self.updateFitness()
        self.updateMetric()

    # This gets overwritten when child classes define their own fitness function
    def fitnessFunc(self, particleIndex):
        p = self.position[particleIndex]
        x = p[0]
        y = p[1]
        return math.sin(math.sqrt((x-math.pi)**2+(y-math.pi)**2)) / math.sqrt((x-math.pi)**2+(y-math.pi)**2)

    def updateFitness(self):
        self.fitness = numpy.zeros(self.nParticles, 'd')
        for i in xrange(self.nParticles):
            self.fitness[i] = self.fitnessFunc(i)
        self.sorted = numpy.argsort(self.fitness).tolist()
        self.sorted.reverse()
        return self.fitnessFunc(self.sorted[0])
        
    def updatePosition(self):
        position = self.position+self.velocity
        position = numpy.minimum(position, self.MAX_POSITION)
        position = numpy.maximum(position,self.MIN_POSITION)
        self.position = position

    def updateVelocity(self):
        best = self.position[self.sorted[0]]
        for dimension in xrange(self.nDimensions):
            for pindex in xrange(self.nParticles):
                self.velocity[pindex][dimension] = \
                            random.uniform(0, self.velocity[pindex][dimension]) + \
                            self.LEARN_RATE*random.uniform(0, \
                            best[dimension]-self.position[pindex][dimension])

    def updateMetric(self):
        self.position = numpy.array(self.position)
        n,m = self.position.shape
        delta = numpy.zeros((n,n),'d')
        for d in xrange(m):
            data = self.position[:,d]
            delta += (data - data[:,numpy.newaxis])**2
        self.metric = numpy.sqrt(delta)

    def plotPoints(self, gp):
        #####
        #This function is hacked to pieces. Clean up.
        #####
        import Gnuplot
        corners = [[0,0],[0,self.spaceSize],[self.spaceSize,0],[self.spaceSize,self.spaceSize]]
        if gp==0:
            gp = Gnuplot.Gnuplot(persist = 0)
            gp('set data style points')
            corners = Gnuplot.PlotItems.Data(corners, with="points 2", title=None)  # No title
            gp.plot(corners)
        else:
            data = [[p[0], p[1]] for p in self.position]
            plot = Gnuplot.PlotItems.Data(data, with="points 3", title=None)  # No title
            gp.plot(plot,corners)
        time.sleep(0.3)
        return gp

    def printData(self):
        #this is just development/debug code
        print 'position\n', self.position
        print
        print 'velocity\n', self.velocity
        print
        print 'metric\n', self.metric
        print
        print 'fitness\n', self.fitness
        print
        print 'sort\n', self.sorted
        print
        print 'best\n', self.fitnessFunc(self.sorted[0])
        print self.position[self.sorted[0]]
#====================================

def main():
    s = Swarm(nDimensions=2,spaceSize=2*math.pi,nParticles=12,learn_rate=.9)
    MAX_ERROR = 0.000001
    errorFunc = lambda :1 - s.updateFitness()
    error = errorFunc()
    print error
    gp = s.plotPoints(0)
    while error>=MAX_ERROR:
        s.updateVelocity()
        s.updatePosition()
        gp = s.plotPoints(gp)
        error = errorFunc()
        print error
    print error
    s.printData()

if __name__ == '__main__':
    for count in xrange(10):
        main()
        time.sleep(1.5)
